import torch
import numpy as np
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets


# 超参数设置
batch_size = 64 # 每一批送入的样本数量
learning_rate = 0.01 # 学习率
momentum = 0.5 # 冲量
EPOCH = 10 # 训练轮数


# 设定一套操作顺序，这个函数将不同的操作组合在一起，之后调用的时候自动对数据进行处理
# transforms.ToTensor()将PILImage格式获numpy.array格式数据转化为张量形式
# 该操作会自动将数据除以255
# 输入模式为（L、LA、P、I、F、RGB、YCbCr、RGBA、CMYK、1）的PIL Image 或 numpy.ndarray (形状为H x W x C)数据范围是[0, 255] 到一个 Torch.FloatTensor，其形状 (C x H x W) 在 [0.0, 1.0] 范围内
# transforms.Normalize对数据进行归一化，需要把传入数据每个通道的值都整理到[-1,1]
# 第一个参数传入平均值mean，第二个传入标准差std，这里对数据 x = (x - mean)/std，
# 0.1307和0.3081是mnist数据集的均值和标准差，
transform = transforms.Compose(
    [
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ]
)

# 加载数据集和测试集，加载或下载位置，是否为训练集，是否需要下载，数据变换的处理
train_dataset = datasets.MNIST(root='./datasets/MNIST', train=True, download=True, transform=transform)  
test_dataset = datasets.MNIST(root='./datasets/MNIST', train=False, download=True, transform=transform)  # train=True训练集，=False测试集

# 自动将输入的数据集按照batch_size进行分割，储存为多个张量，shuffle是否表示需要在每个epoch的时候打乱顺序乱序，true表示打乱，默认false
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


# 显示一部分图片
fig = plt.figure()
for i in range(12):
    plt.subplot(3, 4, i+1)
    plt.tight_layout()
    plt.imshow(train_dataset.train_data[i], cmap='gray', interpolation='none')
    plt.title("Labels: {}".format(train_dataset.train_labels[i]))
    plt.xticks([])
    plt.yticks([])
plt.show()

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 该例子中，输入的是一个单通道的图像, 高宽均为28，维度为batch*1*28*28
        self.conv1 = torch.nn.Sequential(
            # 通道数从1->10,维度为 batch*10*(height - (kernel_size1 - 1))*(weight - (kernel_size1 - 1))
            # 在该例子中，经过Conv2d卷积之后，维度变为了 batch*10*24*24
            torch.nn.Conv2d(1, 10, kernel_size=5), 
            torch.nn.ReLU(), # 激活函数
            # 在该例子中，经过池化后，维度变成了 batch*10*12*12
            torch.nn.MaxPool2d(kernel_size=2), # 通道数不变，高宽变为一半，维度为 batch*10*((height - (kernel_size1 - 1))/2)*((weight - (kernel_size1 - 1)/2)
        )
        self.conv2 = torch.nn.Sequential(
            # 通道数从10->20，维度 batch*20*((height - (kernel_size1 - 1))/2) - (kernel_size2 - 1)*((weight - (kernel_size1 - 1))/2) - (kernel_size2 - 1)
            # 在该例子中，再次经过卷积后，维度变成了 batch*20*8*8
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(), # 激活函数
            # 通道数不变，高宽再变为一半，维度为 batch*20*(((height - (kernel_size1 - 1))/2) - (kernel_size2 - 1))/2*(((weight - (kernel_size1 - 1))/2) - (kernel_size2 - 1))/2
            # 在该例子中，再次经过池化后，维度变成了 batch*20*4*4
            torch.nn.MaxPool2d(kernel_size=2), 
        )
        self.fc = torch.nn.Sequential(
            # 这一步输入的维度为(batch,320)，经过下面两步后，维度变成(batch,10)
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),
        )
        
    def forward(self, x):
        batch_size = x.size(0)
        # 一层卷积层,一层池化层,一层激活层(图是先卷积后激活再池化，差别不大)
        x = self.conv1(x)
        # 再来一次
        x = self.conv2(x)
        # view函数的操作是把矩阵自动变换为某个维度，比如view(x,y)这里将矩阵自动变换为x*y，-1表示自适应
        # flatten 变成全连接网络需要的输入 (batch,20,4,4) ==> (batch,320), 因为经过卷积之后维度为batch*20*4*4
        # 所以这里-1自适应为320,并且，矩阵维度降低到2维
        x = x.view(batch_size, -1)
        # 获取输出
        x = self.fc(x)
        return x
    
model = Net()
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失
# 随机梯度下降法，momentum冲量0-1之间，>0的时候，算法会考虑之前的梯度，有利于收敛
# lr 学习率，每次更新的步长，如果太大会震荡，太小收敛慢
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)  # lr学习率，momentum冲量

# 训练
def train(epoch):
    running_loss = 0.0
    running_total = 0
    running_correct = 0
    
    # batch_idx获取索引，data具体的数据，0表示索引从0开始
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data
        # 梯度归零
        optimizer.zero_grad()
        # 这一步调用了nn.Module中的__call__ : Callable[..., Any] = _call_impl
        # 执行_call_impl函数，会依次执行forward_pre_hook,forward,forward_hook,backward_hook多个函数
        # 其中包括forward，并且这里不能直接写model.forward
        # 这里返回了一个2维张量，维度为batch*10，共有10列，每一列代表归属某一类的预测得分
        outputs = model(inputs)
        # 计算损失，返回的结果是标量(0维张量)
        loss = criterion(outputs, target)
        # 反向传播，反向遍历计算图，链式法则完成，获得每个参数的梯度值
        loss.backward()
        # 通过梯度下降更新参数
        optimizer.step()
        # item表示获取具体的数值
        running_loss += loss.item()
        # dim=1表示输出所在行的最大值，dim=0表示输出所在列的最大值
        # outputs.data表示获取输出张量的原始数据（即张量本身），不涉及梯度，但是在新版本的torch中，可以直接使用outputs，.data操作已经给过时
        _, predicted = torch.max(outputs.data, dim = 1)
        # 总预测数量，输入张量的行的数量，这里感觉是batch_size
        running_total += inputs.shape[0]
        # 预测正确的数量和
        running_correct += (predicted == target).sum().item()
        
        # 每300次计算一次平均损失和准确率
        if batch_idx % 300 == 299:
            print('[%d, %5d]: loss: %.3f , acc: %.2f %%'
                  % (epoch + 1, batch_idx + 1, running_loss / 300, 100 * running_correct / running_total))
            running_loss = 0.0
            running_total = 0
            running_correct = 0
        
# 测试
def test():
    correct = 0
    total = 0
    # 不用计算梯度
    # with语句表示在运行过程中不管是否异常都执行清理操作，释放资源
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    acc = correct / total
    print('[%d / %d]: Accuracy on test set: %.1f %% ' % (epoch+1, EPOCH, 100 * acc))  # 求测试的准确率，正确数/总数
    return acc


# 主函数
if __name__ == "__main__":
    acc_list_test = []
    for epoch in range(0, EPOCH):
        train(epoch)
        acc_test = test()
        acc_list_test.append(acc_test)
    plt.plot(acc_list_test)
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy On TestSet')
    plt.show()